/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.sdk.util;



import com.bff.gaia.unified.sdk.coders.Coder;

import com.bff.gaia.unified.sdk.transforms.DoFn;

import com.bff.gaia.unified.sdk.transforms.DoFnSchemaInformation;

import com.bff.gaia.unified.sdk.values.PCollectionView;

import com.bff.gaia.unified.sdk.values.TupleTag;

import com.bff.gaia.unified.sdk.values.WindowingStrategy;



import java.io.Serializable;

import java.util.Collections;

import java.util.Map;



/**

 * Wrapper class holding the necessary information to serialize a {@link DoFn}.

 *

 * @param <InputT> the type of the (main) input elements of the {@link DoFn}

 * @param <OutputT> the type of the (main) output elements of the {@link DoFn}

 */

public class DoFnInfo<InputT, OutputT> implements Serializable {

  private final DoFn<InputT, OutputT> doFn;

  private final WindowingStrategy<?, ?> windowingStrategy;

  private final Iterable<PCollectionView<?>> sideInputViews;

  private final Coder<InputT> inputCoder;

  Map<TupleTag<?>, Coder<?>> outputCoders;

  private final TupleTag<OutputT> mainOutput;

  private final DoFnSchemaInformation doFnSchemaInformation;



  /**

   * Creates a {@link DoFnInfo} for the given {@link DoFn}.

   *

   * <p>This method exists for backwards compatibility with the Dataflow runner. Once the Dataflow

   * runner has been updated to use the new constructor, remove this one.

   */

  public static <InputT, OutputT> DoFnInfo<InputT, OutputT> forFn(

      DoFn<InputT, OutputT> doFn,

      WindowingStrategy<?, ?> windowingStrategy,

      Iterable<PCollectionView<?>> sideInputViews,

      Coder<InputT> inputCoder,

      TupleTag<OutputT> mainOutput,

      DoFnSchemaInformation doFnSchemaInformation) {

    return new DoFnInfo<>(

        doFn,

        windowingStrategy,

        sideInputViews,

        inputCoder,

        Collections.emptyMap(),

        mainOutput,

        doFnSchemaInformation);

  }



  /** Creates a {@link DoFnInfo} for the given {@link DoFn}. */

  public static <InputT, OutputT> DoFnInfo<InputT, OutputT> forFn(

      DoFn<InputT, OutputT> doFn,

      WindowingStrategy<?, ?> windowingStrategy,

      Iterable<PCollectionView<?>> sideInputViews,

      Coder<InputT> inputCoder,

      Map<TupleTag<?>, Coder<?>> outputCoders,

      TupleTag<OutputT> mainOutput,

      DoFnSchemaInformation doFnSchemaInformation) {

    return new DoFnInfo<>(

        doFn,

        windowingStrategy,

        sideInputViews,

        inputCoder,

        outputCoders,

        mainOutput,

        doFnSchemaInformation);

  }



  public DoFnInfo<InputT, OutputT> withFn(DoFn<InputT, OutputT> newFn) {

    return DoFnInfo.forFn(

        newFn,

        windowingStrategy,

        sideInputViews,

        inputCoder,

        outputCoders,

        mainOutput,

        doFnSchemaInformation);

  }



  private DoFnInfo(

      DoFn<InputT, OutputT> doFn,

      WindowingStrategy<?, ?> windowingStrategy,

      Iterable<PCollectionView<?>> sideInputViews,

      Coder<InputT> inputCoder,

      Map<TupleTag<?>, Coder<?>> outputCoders,

      TupleTag<OutputT> mainOutput,

      DoFnSchemaInformation doFnSchemaInformation) {

    this.doFn = doFn;

    this.windowingStrategy = windowingStrategy;

    this.sideInputViews = sideInputViews;

    this.inputCoder = inputCoder;

    this.outputCoders = outputCoders;

    this.mainOutput = mainOutput;

    this.doFnSchemaInformation = doFnSchemaInformation;

  }



  /** Returns the embedded function. */

  public DoFn<InputT, OutputT> getDoFn() {

    return doFn;

  }



  public WindowingStrategy<?, ?> getWindowingStrategy() {

    return windowingStrategy;

  }



  public Iterable<PCollectionView<?>> getSideInputViews() {

    return sideInputViews;

  }



  public Coder<InputT> getInputCoder() {

    return inputCoder;

  }



  public Map<TupleTag<?>, Coder<?>> getOutputCoders() {

    return outputCoders;

  }



  public TupleTag<OutputT> getMainOutput() {

    return mainOutput;

  }



  public DoFnSchemaInformation getDoFnSchemaInformation() {

    return doFnSchemaInformation;

  }

}